import time
from glob import glob
from multiprocessing import Process, Manager
from multiprocessing import freeze_support

import pandas as pd
import psutil
from tqdm import tqdm
from config import config


# 等待方法
def wait_process():
    available_memory_percent = 100 - psutil.virtual_memory().percent
    if available_memory_percent < 10:
        while True:
            time.sleep(10)
            if available_memory_percent > 20:
                break


def process_data(signal_dict, data_seq, data_list):
    batch_size = signal_dict.get("batch_size")
    data_seq_dict = dict()
    for one in tqdm(range(0, len(data_list), batch_size)):
        # 检测等待
        wait_process()
        # 检测结束

        # 处理数据
        two = one + batch_size
        # 数据读取或者对齐
        one_data = get_data(data_list[one:two])
        data_seq_dict.update(one_data)

    data_seq.append(data_seq_dict)


# 数据处理
def get_data(data_path):
    total_data = dict()
    for one in data_path:
        with open(one, "r", encoding="utf-8") as f:
            one_data = f.read()
        for i in "".join(one_data.split()):
            total_data[i] = ""

    return total_data


# 控制程序
def main_data():
    path_list = glob(config["data_set_path"])

    num_works = 10
    batch_size = 128

    # 参数设置
    signal_dict = Manager().dict()
    data_seq = Manager().list()

    signal_dict["batch_size"] = batch_size

    # 开启异步处理数据
    p_list = []
    for i in range(0, len(path_list), len(path_list) // num_works):
        j = i + len(path_list) // num_works
        p = Process(target=process_data, args=(signal_dict, data_seq, path_list[i:j]))
        p.start()
        p_list.append(p)

    # 训练结束回收进程
    for p in p_list:
        p.join()
    total_dict = dict()
    for one_dict in data_seq:
        for k in one_dict:
            total_dict[k] = one_dict[k]
    pd.to_pickle(["<|pad|>"] + [i for i in total_dict.keys()], config["voc_data_path"])


if __name__ == '__main__':
    freeze_support()

    main_data()
